#include "Common.h"

#define GetUInt(b1, b2, b3, b4) (UINT)(b1 | (b2 << 8) | (b3 << 16) | (b4 << 24))

namespace Process
{
	static BOOL fPrivsAdjusted = FALSE;
	void AdjustPrivileges()
	{
		HANDLE hMyToken;
		HANDLE hThis = GetCurrentProcess();
		if(OpenProcessToken(hThis, TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hMyToken) == FALSE)
		{
			MessageBox(0, "OpenProcessToken", "Error", MB_OK);
			return;
		}
		LUID sLuid;
		if(LookupPrivilegeValue(0, SE_DEBUG_NAME, &sLuid) == FALSE)
		{
			MessageBox(0, "LookupPrivilegeValue", "Error", MB_OK);
			return;
		}
		TOKEN_PRIVILEGES sTokenPriv, sTokenPrivLast;
		sTokenPriv.PrivilegeCount = 1;
		sTokenPriv.Privileges[0].Luid = sLuid;
		sTokenPriv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
		DWORD dwRetLen = 0;
		if(AdjustTokenPrivileges(hMyToken, 0, &sTokenPriv, sizeof(sTokenPriv), &sTokenPrivLast, &dwRetLen) == FALSE)
		{
			MessageBox(0, "AdjustTokenPrivileges", "Error", MB_OK);
			return;
		}
		CloseHandle(hMyToken);
		CloseHandle(hThis);
	}

	HANDLE* GetProcessesByName(__in LPSTR szProcname, __out DWORD* lpdwProcCount, __in bool fCaseSensitive)
	{
		FILE* pLogFile = fopen("GetProcessesByName.txt", "a");
		fprintf(pLogFile, "Searching for process: %s\n", szProcname);
		if(lpdwProcCount == 0)
			return 0;

		*lpdwProcCount = 0;
		DWORD dwProcessList[1024] =  { 0 };
		DWORD dwProcesses = 0;
		Process::AdjustPrivileges();
		if(FALSE == EnumProcesses(dwProcessList, 1024 * sizeof(DWORD), &dwProcesses))
		{
			*lpdwProcCount = 0xFFFFFFFF;	// 0xFFFFFFFF will mean "could not load processes!" 
			return 0;
		}
		DWORD dwProcCount = dwProcesses / sizeof(DWORD);
		HANDLE* hProcesses = (HANDLE*)malloc(sizeof(HANDLE) * dwProcCount);
		fprintf(pLogFile, "Processcount: %u\n", dwProcCount);
		for(DWORD i = 0; i < dwProcCount; ++i)
		{
			DWORD dwProcId = dwProcessList[i];
			fprintf(pLogFile, "Process: %u\n", dwProcId);
			HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS | PROCESS_VM_READ, TRUE, dwProcId);
			if(hProcess == 0)
			{
				fprintf(pLogFile, "\tCould not load process (error: %u)\n", GetLastError());
				continue;
			}
			fprintf(pLogFile, "\tPointer: %p\n", hProcess);
			HMODULE hTmpModules[1] = { 0 };
			DWORD dwModuleSize = 0;
			if(EnumProcessModules(hProcess, hTmpModules, 1 * sizeof(HMODULE), &dwModuleSize) == FALSE)
			{
				CloseHandle(hProcess);
				continue;
			}
			DWORD dwModuleCount = dwModuleSize / sizeof(HMODULE);
			fprintf(pLogFile, "\tModulecount: %u\n", dwModuleCount);
			if(dwModuleCount == 0)
			{
				CloseHandle(hProcess);
				continue;
			}
			BOOL fUseProc = FALSE;
			for(DWORD j = 0; j < 1; ++j)
			{
				HMODULE hCurMod = hTmpModules[j];
				CHAR szFileName[MAX_PATH] = { 0 };
				DWORD dwFileNameLen = GetModuleFileNameEx(hProcess, hCurMod, szFileName, 255);
				if(dwFileNameLen == 0)
					continue;
				fprintf(pLogFile, "\tModuleFileName: %s\n", szFileName);
				CHAR* szFile = strrchr(szFileName, '\\');
				if(szFile == 0)
					continue;
				++szFile;
				fprintf(pLogFile, "\tChrFileName: %s\n", szFile);
				if(fCaseSensitive == FALSE)
				{
					if(_stricmp(szFile, szProcname) == 0)
					{
						hProcesses[(*lpdwProcCount)++] = hProcess;
						fUseProc = TRUE;
						break;
					}
				}
				else
				{
					if(strcmp(szFile, szProcname) == 0)
					{
						hProcesses[(*lpdwProcCount)++] = hProcess;
						fUseProc = TRUE;
						break;
					}
				}
			}
			if(fUseProc == FALSE)
				CloseHandle(hProcess);
		}
		fclose(pLogFile);
		if(*lpdwProcCount != 0)
		{
			HANDLE* hRet = (HANDLE*)malloc(sizeof(HANDLE) * (*lpdwProcCount));
			memcpy(hRet, hProcesses, sizeof(HANDLE) * (*lpdwProcCount));
			free(hProcesses);
			return hRet;
		}
		else
		{
			free(hProcesses);
			return 0;
		}
	}

	HANDLE GetProcessByWindowTitle(__in LPSTR szWindowTitle)
	{
		HWND hwWindow = FindWindow(0, szWindowTitle);
		if(hwWindow == 0)
			return INVALID_HANDLE_VALUE;
		DWORD dwProcId = 0;
		GetWindowThreadProcessId(hwWindow, &dwProcId);
		Process::AdjustPrivileges();
		HANDLE hProc = OpenProcess(PROCESS_ALL_ACCESS | PROCESS_VM_READ, FALSE, dwProcId);
		return hProc == 0 ? INVALID_HANDLE_VALUE : hProc;
	}

	HANDLE GetProcessByClassName(__in LPSTR szClassName)
	{
		HWND hwWindow = FindWindow(szClassName, 0);
		if(hwWindow == 0)
			return INVALID_HANDLE_VALUE;
		DWORD dwProcId = 0;
		GetWindowThreadProcessId(hwWindow, &dwProcId);
		Process::AdjustPrivileges();
		HANDLE hProc = OpenProcess(PROCESS_ALL_ACCESS | PROCESS_VM_READ, FALSE, dwProcId);
		return hProc == 0 ? INVALID_HANDLE_VALUE : hProc;
	}

	void GetProcessFileName(__in HANDLE hProcess, __out CHAR** lpszDest)
	{
		HMODULE hMod[1] = { 0 };
		DWORD dwNeeded = 0;
		if(EnumProcessModules(hProcess, hMod, sizeof(HMODULE), &dwNeeded) == FALSE)
		{
			**lpszDest = 0;
			return;
		}
		if(GetModuleFileNameEx(hProcess, hMod[0], *lpszDest, MAX_PATH) == FALSE)
		{
			**lpszDest = 0;
			return;
		}
	}
};

CMemoryReader::CMemoryReader()
{
	hProcess = INVALID_HANDLE_VALUE;
	szFileName = 0;
	fInitialized = FALSE;
}

CMemoryReader::CMemoryReader(HANDLE hProc)
{
	hProcess = hProc;
	dwProcId = GetProcessId(hProcess);
	szFileName = (LPSTR)malloc(MAX_PATH);
	Process::GetProcessFileName(hProcess, &szFileName);
	fInitialized = TRUE;
	dwLastError = MEMREAD_NO_ERROR;
}

CMemoryReader::~CMemoryReader()
{
	if(szFileName)
		free(szFileName);
}

BOOL CMemoryReader::LoadFromProcessName(__in LPSTR szProcName)
{
	if(fInitialized == TRUE)
	{
		dwLastError = MEMREAD_LOADED_PREV;
		return FALSE;
	}
	DWORD dwProcessCount = 0;
	HANDLE* hProcesses = Process::GetProcessesByName(szProcName, &dwProcessCount, false);
	if(hProcesses == 0 || dwProcessCount == 0)
		return FALSE;
	hProcess = hProcesses[0];
	dwProcId = GetProcessId(hProcess);
	szFileName = (LPSTR)malloc(MAX_PATH);
	Process::GetProcessFileName(hProcess, &szFileName);
	fInitialized = TRUE;
	dwLastError = MEMREAD_NO_ERROR;
	return TRUE;
}

BOOL CMemoryReader::LoadFromWindowName(__in LPSTR szWindowName)
{
	if(fInitialized == TRUE)
	{
		dwLastError = MEMREAD_LOADED_PREV;
		return FALSE;
	}
	HANDLE hProcess = Process::GetProcessByWindowTitle(szWindowName);
	if(hProcess == INVALID_HANDLE_VALUE)
		return FALSE;
	this->hProcess = hProcess;
	dwProcId = GetProcessId(hProcess);
	szFileName = (LPSTR)malloc(MAX_PATH);
	Process::GetProcessFileName(hProcess, &szFileName);
	fInitialized = TRUE;
	dwLastError = MEMREAD_NO_ERROR;
	return TRUE;
}

BOOL CMemoryReader::LoadFromClassName(__in LPSTR szClassName)
{
	if(fInitialized == TRUE)
	{
		dwLastError = MEMREAD_LOADED_PREV;
		return FALSE;
	}
	HANDLE hProcess = Process::GetProcessByClassName(szClassName);
	if(hProcess == INVALID_HANDLE_VALUE)
		return FALSE;
	this->hProcess = hProcess;
	dwProcId = GetProcessId(hProcess);
	szFileName = (LPSTR)malloc(MAX_PATH);
	Process::GetProcessFileName(hProcess, &szFileName);
	fInitialized = TRUE;
	dwLastError = MEMREAD_NO_ERROR;
	return TRUE;
}

UINT CMemoryReader::ReadUInt(DWORD dwAdress)
{
	if(!fInitialized == TRUE)
	{
		dwLastError = MEMREAD_NOT_LOADED;
		return ~((UINT)0);
	}
	dwLastError = MEMREAD_NO_ERROR;
	return Read<UINT>(dwAdress);
}

ULONGLONG CMemoryReader::ReadUInt64(DWORD dwAdress)
{
	if(!fInitialized == TRUE)
	{
		dwLastError = MEMREAD_NOT_LOADED;
		return ~((ULONGLONG)0);
	}
	dwLastError = MEMREAD_NO_ERROR;
	return Read<ULONGLONG>(dwAdress);
}

FLOAT CMemoryReader::ReadFloat(DWORD dwAdress)
{
	if(!fInitialized == TRUE)
	{
		dwLastError = MEMREAD_NOT_LOADED;
		return 0xFFFFFFFF;
	}
	dwLastError = MEMREAD_NO_ERROR;
	return Read<FLOAT>(dwAdress, TRUE);
}

INT CMemoryReader::ReadInt(DWORD dwAdress)
{
	if(!fInitialized == TRUE)
	{
		dwLastError = MEMREAD_NOT_LOADED;
		return ~((INT)0);
	}
	dwLastError = MEMREAD_NO_ERROR;
	return Read<INT>(dwAdress);
}

BOOL CMemoryReader::ReadString(DWORD dwAdress, LPSTR lpszDest, DWORD dwMaxLen)
{
	LPSTR szMem = (LPSTR)dwAdress;
	UINT i = 0;
	for( ; i < dwMaxLen - 1; ++i)
	{
		if(*(szMem + i) != 0)
		{
			*(lpszDest + i) = *(szMem + i);
		}
		else
			break;
	}
	*(lpszDest + i) = 0;
	return TRUE;
}

BYTE* CMemoryReader::ReadBytes(DWORD dwAdress, DWORD dwLen)
{
	DWORD dwRead;
	BYTE* out = new BYTE[dwLen];
	ReadProcessMemory(hProcess, (LPCVOID)dwAdress, out, dwLen, &dwRead);
	return out;
}

template<typename T>
T CMemoryReader::Read(unsigned dwAdress, BOOL fFixRet)
{
	Process::AdjustPrivileges();
	T gOut = 0;
	DWORD dwBytesRead = 0;
	if(ReadProcessMemory(hProcess, (LPCVOID)dwAdress, &gOut, sizeof(T), &dwBytesRead) == FALSE)
	{
		gOut = 0xFFFFFFFF;
		dwLastError = MEMREAD_READ_FAILED;
		/*char msg[255];
		sprintf(msg, "Error: %u - Bytes Read: %u - Bytes to read: %u - Adress: %X - Process: %p", GetLastError(), dwBytesRead, sizeof(T), dwAdress, hProcess);
		MessageBox(0, msg, "", MB_OK);*/
	}
	return gOut;
}

void CMemoryReader::SetErrorStream(std::ostream& ostr)
{
	//m_errorstream = ostr;
}